import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
# read mat file
# from scipy.io import loadmat
import h5py as hp
# split dataset
from sklearn.model_selection import train_test_split
# construct dataset
import torch.utils.data as Data
import matplotlib.pyplot as plt
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def NN(Path, data_base):
    # Generator
    class Generator(nn.Module):
        def __init__(self):
            super(Generator, self).__init__()
            self.Generator_net = nn.Sequential(
                nn.Linear(99, 288),
                nn.BatchNorm1d(288),
                nn.LeakyReLU(),

                nn.Linear(288, 256),
                nn.BatchNorm1d(256),
                nn.LeakyReLU(),

                nn.Linear(256, 64),
                nn.BatchNorm1d(64),
                nn.LeakyReLU(),
                nn.Linear(64, 8),
            )

        def forward(self, x):
            y = self.Generator_net(x)
            return y

    G_net = Generator()
    G_net.load_state_dict(torch.load(Path))

    G_net.eval()

    with torch.no_grad():
        # 导入测试集
        y = 0
        y_hat = 0
        Index = 0
        for Z_Condition, real_data in data_base:
            # G net output
            y_G = G_net(Z_Condition).reshape(-1)
            # real data
            real = real_data.reshape(-1)
            # NMSE
            if Index == 0:
                y = real
                y_hat = y_G
                Index = Index + 1
            else:
                y = torch.cat((y, real), 0)
                y_hat = torch.cat((y_hat, y_G), 0)

    import scipy.io as so
    file = 'G_Net_' + str(20) + 'GHZ_' + str(50) + 'km_' + str(1) + 'Span_' + str(10) + 'dBm.mat'
    so.savemat(file, {'TxSignal_X': y_hat.cpu()})
